import pytorch_lightning as pl
from monai.data import DataLoader
from monai.transforms import (
    Compose,
    EnsureChannelFirstd,
    EnsureTyped,
    RandFlipd,
    RandRotate90d,
    RandSpatialCropd,
    ScaleIntensityd,
)
from datasets import DATASETS


class CellDataModule(pl.LightningDataModule):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.batch_size = config['data']['batch_size']
        
        # 定义数据增强和预处理流程
        self.train_transforms = Compose([
            ScaleIntensityd(keys=["image"]),
            RandSpatialCropd(keys=["image", "label"], roi_size=(128, 128), random_size=False),
            RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=0),
            RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=1),
            RandRotate90d(keys=["image", "label"], prob=0.5, max_k=3),
            EnsureTyped(keys=["image", "label"]),
        ])
        
        self.val_transforms = Compose([
            ScaleIntensityd(keys=["image"]),
            EnsureTyped(keys=["image", "label"]),
        ])
        
        # 使用注册器创建数据集
        dataset_class = DATASETS.get(config['data']['dataset_type'])
        dataset_args = config['data'].get('dataset_args', {})
        
        self.train_ds = dataset_class(
            transform=self.train_transforms,
            **dataset_args
        )
        self.val_ds = dataset_class(
            transform=self.val_transforms,
            **dataset_args
        )
        
    def train_dataloader(self):
        return DataLoader(
            self.train_ds, 
            batch_size=self.batch_size, 
            shuffle=True, 
            num_workers=self.config['data']['num_workers']
        )
        
    def val_dataloader(self):
        return DataLoader(
            self.val_ds, 
            batch_size=self.batch_size, 
            num_workers=self.config['data']['num_workers']
        )
